/** \file    kDimTree.h
  * \brief   provides quick search of points using K-dimensional Tree
  * \date    1998-2001
  * \author  Aloudin (Albert) Akhriev
  * \version $Header: /home/cscom1r4/cvsroot/cs/include/xtd/kDimTree.h,v 1.2 2004/11/22 19:46:41 cscom1r4 Exp $
  */

#ifndef _kDimTree_h_
#define _kDimTree_h_

#include <vector>
#include <cassert>

namespace xtd {

//==========================================================================
// K-dimensional Tree.
//
// Description sources:
// 1. Z.Zhang, "Iterative Point Matching for Registration of Free-Form
//    Curves and Surfaces," IJCV, 13:2, 119-152, 1994.
// 2. F.Preparata, and M.Shamos, "Computational Geometry, An Introduction,"
//    Springer, Berlin, Heidelberg, New-York, 1986.
//
// 'TYPE'  - type of qPoint, must be a kind of vector type with operator[].
//           E.g.: TYPE = Vec2f.
//			 (method 'TYPE::Size()' is abolished by Fedor since it was used only inside assert)
// 'QTYPE' - query type. For the qPoint of 'QTYPE' type the nearest one of type
//           'TYPE' should be found on the tree. Usually 'QTYPE' == 'TYPE'.
//           User should define method 'Distance(TYPE&)' for 'QTYPE' points.
// 'DIST'  - type of distance between points. E.g. DIST = double.
// 'DIM'   - dimension of tree (dimension of vector-qPoint, TYPE.Size() == DIM).
//
// Problems:
// 1. Coordinates of qPoint may be of integer type, but distance between points
//    is of float type. I use DIST type in all cases (see Search(..)), but it
//    will be better to introduce separate coordinates type for efficiency.
//==========================================================================

///automatic Distance() provider
template <class TYPE>
class QTYPE: public TYPE {
public:
	QTYPE() {}
	QTYPE(const TYPE &s) : TYPE(s) {}
	double Distance(const TYPE &s) const
		{ return ::abs(*this - s); }
};

template<class TYPE, class QTYPE, class DIST, int DIM>
class kDimTree
{
public:
  /**\brief Function recursively creates tree from the array of points. Generated
   * tree is optimally balanced. Ok == true. */
  bool Create(const TYPE * pPoints, int length);
  ///convinient helper: creates temporary array of TYPE for you, *I should be convertable into TYPE
  template <class I, class EI>
  bool Create(I begin, EI end, int length);

  /**\brief Function recursively searches tree for nearest qPoint to the given one,
   * using maximal displacement on qPoint's location ('distThreshold').
   * The function never returns point with zero distance to qPoint if 'selfIgnore' is true.
   * Returns: Ok == pointer to the nearest qPoint, otherwise NULL. */
  const TYPE * FindNearestPoint(const QTYPE & qPoint, DIST distThreshold, 
	  bool selfIgnore=false, bool ignore=false, QTYPE ignorePt = QTYPE()) const;

  /**\brief Function recursively searches tree for the nearest neighbors of given
   * qPoint. Function uses maximal displacement on neighbor's location
   * ('upBound'). Returns array of found neighbors. */
  const std::vector<TYPE*> * 
  FindNeighbors(const QTYPE & qPoint, DIST distThreshold) const;

  /// Function clears tree's structure.
  bool Clear();

  // Constructor/Destructor.
  kDimTree() : m_pRoot(NULL), m_pNeiPt(NULL) { Clear(); }
  ~kDimTree() { Clear(); }

  // Info-functions.
  bool IsEmpty() const { return (m_pRoot == NULL); }
  int  NodesNum() const { return m_nNodes; }

public:
  typedef  kDimTree<TYPE,QTYPE,DIST,DIM>*  PTHIS;
  typedef  std::vector<TYPE*>              PointPtrArr;

  struct Node           ///< Node of kD-tree.
  {
    Node * pL;          ///< left child node
    Node * pR;          ///< right child node
    TYPE   qPoint;       ///< qPoint that splits space into two parts
    int    dir;         ///< direction of space splitting (index of coord. axis that is perpendicular to the splitting plane)

    Node() : pL(NULL), pR(NULL), qPoint(), dir(-1) {}
    ~Node() { delete pL;  delete pR; }

	///calls 'f' with each node of subtree as argument
	template <class F>
	void enumSubTree(F &f) const
		{ f(qPoint); 
		  if (pL) pL->enumSubTree(f);
		  if (pR) pR->enumSubTree(f); }
  };

  Node *        m_pRoot;     ///< root node of this tree
  int           m_nNodes;    ///< number of tree's nodes
  TYPE *        m_pBestPt;   ///< best found qPoint (temp. variable)
  DIST          m_upBound;   ///< threshold on distance (temp. variable)
  PointPtrArr * m_pNeiPt;    ///< neighbour points of selected one (temp. var.)
  bool			m_selfIgnore;///< do not return point with zero distance to query point
  bool			m_ignore;	 ///< the same but for m_ignorePt
  QTYPE			m_ignorePt;

public:
  #ifdef _DEBUG
  int           m_nComp;     ///< number of comparisons made (debug)
  #endif

private:
  /**\brief Function sorts an array into ascending numerical order by Shell's method.
   * Parameter 'dir' (direction of space cutting) denotes index of coordinate
   * of object of type 'TYPE', that is used for comparison. */
  void Sort(TYPE ** arr, int length, int dir);

  /**\brief Function recursively splits space into two halfs. Direction of splitting
   * is cyclically changed on each level. Returns newly created node. */
  Node * CreateBranch(TYPE ** pa, int length, const Node * pParent);

  /// Function recursively searches for nearest qPoint to the given one.
  void Search(const Node * pn, const QTYPE & qPoint) const;

  /// Function recursively searches for nearest neighbors of given qPoint.
  void SearchNeighbors(const Node * pn, const QTYPE & qPoint) const;
};

////////////////////
// implementation //
////////////////////

template<class TYPE, class QTYPE, class DIST, int DIM>
void kDimTree<TYPE,QTYPE,DIST,DIM>::Sort(TYPE ** arr, int length, int dir)
{
    for(int gap = length>>1; gap > 0; gap >>= 1)
    {
      for(int i = gap; i < length; i++)
      {
        for(int j = i-gap; j >= 0 && (*(arr[j]))[dir] > (*(arr[j+gap]))[dir];)
        {
          TYPE * p = arr[j];
          arr[j] = arr[j+gap];
          arr[j+gap] = p;
          j -= gap;
        }
      }
    }
}

template<class TYPE, class QTYPE, class DIST, int DIM>
typename kDimTree<TYPE,QTYPE,DIST,DIM>::Node * 
	kDimTree<TYPE,QTYPE,DIST,DIM>::CreateBranch(TYPE ** pa, int length, const Node * pParent)
{
    Node * pn = NULL;
    if (length > 0)
    {
      int half = length >> 1;
      int dir = (pParent != NULL) ? ((pParent->dir+1) % DIM) : 0;
      Sort(pa, length, dir);
      pn = new Node();
      pn->dir = dir;
      pn->qPoint = *(pa[half]);
      if (length > 1)
      {
        pn->pL = CreateBranch(pa, half, pn);
        pn->pR = CreateBranch(pa + (half+1), length - (half+1), pn);
      }
      m_nNodes++;
    }
    return pn;
}

template<class TYPE, class QTYPE, class DIST, int DIM>
void kDimTree<TYPE,QTYPE,DIST,DIM>::Search(const Node * pn, const QTYPE & qPoint) const
{
    DIST c1 = qPoint[ pn->dir ];
    DIST c2 = pn->qPoint[ pn->dir ];
    DIST dc = (c2 >= c1) ? (c2-c1) : (c1-c2);    // DIST may be unsigned

    if (dc < m_upBound)
    {
      DIST t = qPoint.Distance(pn->qPoint);
      if (t < m_upBound && 
		  !(m_selfIgnore && t == 0) && 
		  !(m_ignore && m_ignorePt.Distance(pn->qPoint) == 0))
      {
        ((PTHIS)this)->m_pBestPt = (TYPE*)(&(pn->qPoint));
        ((PTHIS)this)->m_upBound = t;
      }
    }

    #ifdef _DEBUG
    ((PTHIS)this)->m_nComp++;
    #endif

    if (pn->pL != NULL && c1 < (c2+m_upBound))
      Search(pn->pL, qPoint);
    if (pn->pR != NULL && c2 < (c1+m_upBound))
      Search(pn->pR, qPoint);
}

template<class TYPE, class QTYPE, class DIST, int DIM>
void kDimTree<TYPE,QTYPE,DIST,DIM>::SearchNeighbors(const Node * pn, const QTYPE & qPoint) const
{
    DIST c1 = qPoint[ pn->dir ];
    DIST c2 = pn->qPoint[ pn->dir ];
    DIST dc = (c2 >= c1) ? (c2-c1) : (c1-c2);    // DIST may be unsigned

    if (dc <= m_upBound && qPoint.Distance(pn->qPoint) <= m_upBound)
    {
      ((PTHIS)this)->m_pNeiPt->push_back(&(pn->qPoint));
    }

    #ifdef _DEBUG
    ((PTHIS)this)->m_nComp++;
    #endif

//    if (pn->pL == NULL && pn->pR == NULL)
//    return;                                    // leaf

    if (pn->pL != NULL && c1 < c2+m_upBound)
      SearchNeighbors(pn->pL, qPoint);
    if (pn->pR != NULL && c2 < c1+m_upBound)
      SearchNeighbors(pn->pR, qPoint);
}

template<class TYPE, class QTYPE, class DIST, int DIM>
bool kDimTree<TYPE,QTYPE,DIST,DIM>::Create(const TYPE * pPoints, int length)
{
    Clear();
    if (length < 1)
      return false;

//    assert(DIM > 1 && pPoints[0].Size() == DIM);
    TYPE ** pa = new TYPE* [length];
	try {
		for(int k = 0; k < length; k++)
		  pa[k] = (TYPE*)(pPoints++);
		m_pRoot = CreateBranch(pa, length, NULL);
		delete[] pa;
		return true;
	}
	catch(...) {
		delete[] pa;
		throw;
	}
}

template<class TYPE, class QTYPE, class DIST, int DIM>
template <class I, class EI>
bool kDimTree<TYPE,QTYPE,DIST,DIM>::Create(I begin, EI end, int length)
{
    TYPE * p = new TYPE [length];
	try {
		for(int k = 0; k < length && begin != end; ++k, ++begin)
			p[k] = *begin;
		bool res = Create(p, length);
		delete[] p;
		return res;
	}
	catch(...) {
		delete[] p;
		throw;
	}
}

template<class TYPE, class QTYPE, class DIST, int DIM>
const TYPE * kDimTree<TYPE,QTYPE,DIST,DIM>::FindNearestPoint(
	const QTYPE & qPoint, DIST distThreshold, 
	bool selfIgnore, bool ignore, QTYPE ignorePt) const
{
    #ifdef _DEBUG
    ((PTHIS)this)->m_nComp = 0;
    #endif

    ((PTHIS)this)->m_pBestPt = NULL;
    ((PTHIS)this)->m_upBound = distThreshold;
	((PTHIS)this)->m_selfIgnore = selfIgnore;
	((PTHIS)this)->m_ignore = ignore;
	((PTHIS)this)->m_ignorePt = ignorePt;

    if (m_pRoot != NULL)
      Search(m_pRoot, qPoint);
    return m_pBestPt;
}

template<class TYPE, class QTYPE, class DIST, int DIM>
const std::vector<TYPE*> * 
kDimTree<TYPE,QTYPE,DIST,DIM>::FindNeighbors(const QTYPE & qPoint, DIST distThreshold) const
{
    #ifdef _DEBUG
    ((PTHIS)this)->m_nComp = 0;
    #endif

    if (m_pNeiPt == NULL)
      ((PTHIS)this)->m_pNeiPt = new PointPtrArr();
    if (m_pNeiPt->capacity() < m_nNodes)
      ((PTHIS)this)->m_pNeiPt->reserve(m_nNodes);

    ((PTHIS)this)->m_pNeiPt->clear();
    ((PTHIS)this)->m_upBound = distThreshold;
    if (m_pRoot != NULL)
      SearchNeighbors(m_pRoot, qPoint);
    return m_pNeiPt;
}

template<class TYPE, class QTYPE, class DIST, int DIM>
bool kDimTree<TYPE,QTYPE,DIST,DIM>::Clear() {
    #ifdef _DEBUG
    m_nComp = 0;
    #endif

    delete m_pRoot;
    m_pRoot = NULL;
    delete m_pNeiPt;
    m_pNeiPt = NULL;
    m_nNodes = 0;
    m_pBestPt = NULL;
    m_upBound = 0;
    return false;
}

} //namespace xtd;

#endif //_kDimTree_h_
